#!/usr/bin/python
#coding:utf-8

# Copyright 2012 Nicolau Leal Werneck, Anna Helena Reali Costa and
# Universidade de São Paulo
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

###############################################################################
## Open a (pinhole camera model) picture, extract the edgels, and find
## the orientation using SQP.
##

import sys

import time

import matplotlib

if __name__ == '__main__':
    if sys.argv[0][-7:] == '-nox.py':
        matplotlib.use('Agg') 

import plot_aux

from pylab import *

import scipy.io

import scipy.ndimage

from numpy import dot

from camori import PicturePinhole, PictureEqrec, quaternion_to_matrix, measure_error, quaternion_product, dir_colors

from camori_aux import generate_mask

from quaternion import Quat, random_quaternion

import simplejson

import Image

set_printoptions(precision=3)

from pylab import *

from scipy.misc import factorial as fact
from sqp_step1 import create_pic

class EdgelExtractor:
    scharr_filter = np.array([
            [ -3, 0, 3],
            [-10, 0,10],
            [ -3, 0, 3] ])*np.sqrt(1/254.0)

    def zernike_masks(self,N):
        yy,xx = (mgrid[0.:N,:N] - (N/2) ) *2 / N

        return [generate_mask(self.Nw, 2048/self.Nw,0),
                generate_mask(self.Nw, 2048/self.Nw,1),
                generate_mask(self.Nw, 2048/self.Nw,2),
                generate_mask(self.Nw, 2048/self.Nw,3)]


    def __init__(self, frame, Nw = 7, Vnm = []):
        self.Nw = Nw
        if Vnm ==[]:
            self.Vnm = self.zernike_masks(self.Nw)
        else:
            self.Vnm = Vnm
        # for nm in range(4):
        #     self.Vnm[nm] = self.Vnm[nm] / sqrt((self.Vnm[nm]**2).sum())

        self.frame = frame

        self.gradx = scipy.ndimage.convolve(self.frame,
                               EdgelExtractor.scharr_filter)
        self.grady = scipy.ndimage.convolve(self.frame,
                               EdgelExtractor.scharr_filter.T)
        
    def extract_zm_simple(self, ex, ey):
        # ex = np.round(rex)
        # ey = np.round(rey)
        #snp = self.frame[ey-3:ey+4,ex-3:ex+4]
        snp = self.frame[ey-self.Nw/2:ey+self.Nw/2+1,ex-self.Nw/2:ex+self.Nw/2+1]
        A00 = (self.Vnm[0] * snp).sum()
        A11r = (self.Vnm[1] * snp).sum()
        A11i = (self.Vnm[1].T * snp).sum()
        A20 = (self.Vnm[2] * snp).sum()
        A31r = (self.Vnm[3] * snp).sum()
        A31i = (self.Vnm[3].T * snp).sum()

        A11m = sqrt(A11r**2 + A11i**2)
        A31_ = (A31r*A11r + A31i*A11i)/A11m

        ## 1993 article, using A00 A11 A20
        phi = arctan2(A11i, A11r)
        l = A20 / A11m
        k = 3*A11m / (2 * (1-l**2)**1.5)
        h = (A00 - k*pi/2 +k * arcsin(l) + k * l *sqrt(1-l**2))/pi

        l15 = (1-l**2)**1.5

        # a20hat = 2*k*l15*l/3
        # a31hat = -2*k*(1-6*l**2)*l15/15
        # print 70*'-'
        # print 'l:',l,'h:',h
        # print 'A11', A11r, A11i, A11m
        # print 'A20', A20, 'A20_', a20hat
        # print 'A31', A31r,A31i, 'A31_', A31_, 'A31*', a31hat


        nphi=phi
        nl = ((sqrt(25 * A31_**2 + 24 * A20**2) + 5 * A31_) /
             (12 * A20))
        nk = -15 * A31_ / (2*sqrt(1-l**2)*(1-l**2)*(1-6*l**2)) 

        return phi,l,k,h,nphi,nl,nk, A11r, A11i

    def extract_zm(self, ex, ey):
        snp = self.frame[ey-self.Nw/2:ey+self.Nw/2+1,ex-self.Nw/2:ex+self.Nw/2+1]
        # snp = self.frame[ey-3:ey+4,ex-3:ex+4]
        A00 = (self.Vnm[0] * snp).sum()
        A11r = (self.Vnm[1] * snp).sum()
        A11i = (self.Vnm[1].T * snp).sum()
        A20 = (self.Vnm[2] * snp).sum()
        A31r = (self.Vnm[3] * snp).sum()
        A31i = (self.Vnm[3].T * snp).sum()

        A11m = sqrt(A11r**2 + A11i**2)
        A31m = sqrt(A31r**2 + A31i**2)

        A11r += A31r
        A11i += A31i

        #print A11r, A11i, A20, A31r, A31i

        ## First implementation, old article
        # th = arctan2(A11i, A11r)
        # amp = sqrt(A11r*A11r + A11i*A11i)
        # l = -3.5 * A20/amp
        # ux = A11r/amp
        # uy = A11i/amp
        # px = ex + l * ux
        # py = ey + l * uy
        # return amp,px,py,ux,uy,th,l, snp[3,3]


        ## 1993 article, using A00 A11 A20
        # phi = arctan2(A11i, A11r)
        # l = A20 / A11m
        # k = 3*A11m / (2 * (1-l**2)**1.5)
        # h = (A00 - k*pi/2 +k * arcsin(l) + k * l *sqrt(1-l**2))/pi

        ## 1997 article, using A11, A20 and A31
        amp = sqrt(A11r*A11r + A11i*A11i)
        ux = A11r/amp
        uy = A11i/amp

        A31m = sqrt(A31r**2 + A31i**2)

        ## New article, using A31        
        phi = arctan2(A11i, A11r)
        amp = sqrt(A11r*A11r + A11i*A11i)
        ux = A11r/amp
        uy = A11i/amp
        A31m = sqrt(A31r**2 + A31i**2)

        # l = ((sqrt(25 * A31m** A31m + 24 * A20*A20) + 5 * A31m) /
        #     (12 * A20))
        l = 3.5 * A20/amp
        
        h = -15 * A31m / (2*sqrt(1-l**2)*(1-l**2)*(1-6*l**2)) 
        #px = ex + l * ux
        #py = ey + l * uy
        return amp,px,py,ux,uy,phi,l,h,snp[3,3]




    def modelo_zm(self, ex, ey):
        snp = self.frame[ey-self.Nw/2:ey+self.Nw/2+1,ex-self.Nw/2:ex+self.Nw/2+1]
        A00 = (self.Vnm[0] * snp).sum()
        A11r = (self.Vnm[1] * snp).sum()
        A11i = (self.Vnm[1].T * snp).sum()
        A20 = (self.Vnm[2] * snp).sum()
        A31r = (self.Vnm[3] * snp).sum()
        A31i = (self.Vnm[3].T * snp).sum()

        mod = zeros(snp.shape)
        mod += A00 * self.Vnm[0]
        mod += A11r * self.Vnm[1]
        mod += A11i * self.Vnm[1].T
        mod += A20 * self.Vnm[2]
        mod += A31r * self.Vnm[3]
        mod += A31i * self.Vnm[3].T

        return snp, mod




def integrate_windowed_poly(N, M, nm):
    out = zeros((N,N))

    zernike_polys = [
        [1],
        [1,0],
        [2,0,-1],
        [3,0,-1,0],
        [1,0,0,0]   ]



    for j in mgrid[0.:N]:
        for k in mgrid[0.:N]:
            for jm in mgrid[0:M]:
                for km in mgrid[0:M]:
                    xx = 2 * (k*M + km + .5)/(N*M) - 1.0
                    yy = 2 * (j*M + jm + .5)/(N*M) - 1.0
                    rho = sqrt(xx*xx+yy*yy)
                    if rho > 1:
                        vv = 0
                    else:
                        vv = np.polyval(zernike_polys[nm], rho)                        
                    out[j,k] += vv
    return out/(M*M)


def synth_test():
        ## Test on very simple artificial step image patches.
    if False:
        steps = []

        stp = []
        for h in [50,100,150,200]:
            for l in mgrid[0.:5:1]/4:
                st = zeros((7,7))

                # for j in range(7):
                #     for k in range(7):
                #         st[j,k] = h if j<k else 250 if j>k else h*l + (1-l) * 250
                st[:,:3] = h
                st[:,3] = h*l + (1-l) * 250
                st[:,4:] = 250
                steps.append(st)
                stp.append((l-.5,h))

        kk = []

        for n in range(len(steps)):
            frm = steps[n]
            ee = EdgelExtractor(frm, Nw=Nw, Vnm=Vnm)
            phi,l,k,h,nphi,nl,nk = ee.extract_zm_simple(3,3)
            print
            print stp[n][0], 3.5*l, 3.5*nl

            kk.append((l,nl))
        kk = array(kk)





if __name__ == '__main__':
    rc('text',usetex=False)
    set_printoptions(precision=2)
    ion()    

    Nw = 7

    ## Values calculated with Mathematica. Not so good results. :(
    Vnm = np.loadtxt('zernike7x7-mathematica.dat')
    Vnm = Vnm.reshape((4,7,7))
    ## Calculate masks using a crowbar and a lot of patience.
    # Vnm = [generate_mask(Nw, 2048,0),
    #        generate_mask(Nw, 2048,1),
    #        generate_mask(Nw, 2048,2),
    #        generate_mask(Nw, 2048,3)]
    ## Load results from the previous method.
    # Vnm = np.load('zernike_masks7x7.npz')['arr_0']


    ## Avoid zero divide warnins...
    np.seterr(divide='ignore')

    PlotStuff=True
    if sys.argv[0][-7:] == '-nop.py':
        PlotStuff=False

    if PlotStuff:
        ## Plot stuff immediately
        ion()

    #################################################################
    ## Load images and initialize pic objects

    ## Sets filename from input argument
    if len(sys.argv) < 3:
        print sys.argv[0], '<job_file.json> <frame_number>'
        raise Exception('Insufficient number of parameters')

    finput = open(sys.argv[1])
    job_params = simplejson.load(finput)
    finput.close()

    fileroot = job_params['root_directory']

    framenum = int(sys.argv[2])

    pic = create_pic(job_params, framenum)

    gspc = job_params['grid_spacing']
    glim = job_params['gradient_threshold'] # ** -1
    ## Run the grid sweep, and generate the edgels list.
    ## Edgel extractor parameters
    pic.extract_edgels(gspc, glim)



    ## img1
    #19 152
    #326 57
    ## img9
    #190 149
    #634 118


    frm = pic.frame[:,:].sum(2)/3.
    figure(11)
    imshow(frm/256., cmap=cm.bone)


    ee = EdgelExtractor(frm, Nw=7, Vnm=Vnm)

    #px1,px2,py1,py2 = 19,326,152,57 # img 1
    #px1,px2,py1,py2 = 184.,592.,80.,47. # img 2
    #px1,px2,py1,py2 = 190,634,149,118 # img 9    

    #px1,px2,py1,py2 = 280.,419.,66.75,21.8 # img 19

    #px1,px2,py1,py2 = 600.,642.,273.4,282.3 # apa img 0
    #px1,px2,py1,py2 = 305,375.0,204.45,220.65 # apa img 0
    #px1,px2,py1,py2 = 420.0,700.0,736.7,738.3 # apa img 1

    #px1,px2,py1,py2 = 314.0,608.0,71.0,60.0 # ap9a img 1
    #px1,px2,py1,py2 = 414.0,514.0,65.95,62.25 # ap9a img 1

    #px1,px2,py1,py2 = 19,326,151,56 # img 1
    #px1,px2,py1,py2 = 19,113,151.8,122.0 # img 1 trecho
    px1,px2,py1,py2 = 143,234,112.6,84.45 # img 1

    #px1,px2,py1,py2 = 220.0,340.0,77.2,67.6 # img 2 trecho
    #px1,px2,py1,py2 = 220.0,408.0,77.2,61.8 # img 2 trecho
    #px1,px2,py1,py2 = 430.,530.,60.3,52.15 # img 2 trecho


    #px1,px2,py1,py2 = 186,483,148.6,127.6 # img 9 trecho




    # a,b = ee.modelo_zm(px1+6,py1)
    # subplot(1,2,1)
    # imshow(a, cmap=cm.bone, vmin=0,vmax=256)
    # subplot(1,2,2)
    # imshow(b, cmap=cm.bone, vmin=0,vmax=256)
    # myxtplikz()


    uxR = -(py2-py1)
    uyR = (px2-px1)
    nf = 1.0/sqrt(uxR*uxR + uyR*uyR)
    uxR, uyR = nf*uxR, nf*uyR

    # B11r = -uxR
    # B11i = -uyR

    # qq=[]
    # for ex in range(int(px1), int(px2), 1):
    #     ey = py1 + (ex - px1) * (py2 - py1) / (px2-px1)

    #     rex = ex
    #     rey = np.round(ey)

    #     phi,ll,ak,ah,nphi,nll,k,A11r,A11i =  ee.extract_zm_simple(rex,rey)
    #     nf = 1.0/sqrt(A11r*A11r + A11i*A11i)
    #     A11r, A11i = nf*A11r, nf*A11i
    #     err = B11r*A11i-B11i*A11r
    #     #qq.append([3.5* nll, rex])
    #     qq.append([rex, ee.frame[rey,rex]])
        
    #     print 3.5*nll, A11r, A11i, B11r, B11i, err, rex, rey
        
    # qq = array(qq)
    # xx = mgrid[0.:qq.shape[0]]/qq.shape[0] - .5

    # qq[:,1] -= qq[:,1].min()
    # qq[:,1] /= qq[:,1].max()
    # qq[:,1] -=.5

    # figure(1)
    # plot(sort(qq[:,1]), xx, '-+')
    # plot(xx,xx,'k--')
    # grid()
    # #axis('equal')
    # myxtplk()



    if True:
        ppa = []
        pp = []
        pr = []
        vv = array([-1,1])
        ov = array([0,1])
        lL = []

        for ex in range(int(px1), int(px2), 1):
            ## Get coordinates from my line
            ey = py1 + (ex - px1) * (py2 - py1) / (px2-px1)

            rex = ex
            rey = np.round(ey)

            print 70*'-'
            phi,ll,ak,ah,nphi,nll,k,A11r,A11i =  ee.extract_zm_simple(rex,rey)
            al = ll * 3.5
            l = nll*3.5
            
            ux = cos(phi)
            uy = sin(phi)

            px = rex
            py = rey + l*(ux**2/uy + uy)

            plot(rex,rey, 'b+')
            plot(ex,ey, 'r+')
            # plot(px+7*uy*vv, py-7*ux*vv, 'r-')
            plot(rex-6*uxR*ov, rey-6*uyR*ov, 'g-')
            plot(rex+6*ux*ov, rey+6*uy*ov, 'r-')
            plot(rex+l*ux*ov, rey+l*uy*ov, 'y-')
            plot(px, py, 'yx')

            # plot(ex, l/uy + ey - err, 'r+')
            
            # eyR = py1 + (ex - px1) * (py2 - py1) / (0.+px2-px1) #img 9
            ppa.append((arctan2(uy,ux), al, ak, ah))
            pp.append((arctan2(uy,ux), l, k))

            ip = ee.frame[rey,rex]
            
            pr.append((ux*uyR-uy*uxR, py, ey, ip, l*(ux**2/uy + uy), ey-rey, l) )
            ol = ey-rey
            lL. append( (ol , al*(ux**2/uy + uy), l*(ux**2/uy + uy)))
        ## plot the line guess
        plot([px1,px2],[py1,py2], 'g-+')

        lL = array(lL)
        ppa = array(ppa)
        pp = array(pp)
        pr = array(pr)

        Np = pp.shape[0]

        m0 = sort(pp[:,0])[Np/2]
        m1 = sort(pp[:,1])[Np/2]

        xx = mgrid[0.:Np]/Np

        figure(2)
        plot((180/pi)*np.sort(np.abs(pp[:,0]-m0)),xx)
        grid()

        figure(3)
        plot(pr[:,1]-pr[:,2], '-+')
        ylim(1,-1)
        # twinx()
        # plot(pr[:,3],'g-+')
        grid()


        figure(4)
        suptitle('Biased errors on edgel extraction')
        subplot(1,2,1)
        plot(pr[:,6],(180/pi)*pr[:,0], 'o')
        grid()


        ii = argsort(pr[:,6])
        a,b,c=polyfit(pr[:,6],(180/pi)*pr[:,0],2)
        plot(pr[ii,6], a*pr[ii,6]**2+b*pr[ii,6]+c,'-g')

        subplot(1,2,2)
        plot(pr[:,6],pr[:,1]-pr[:,2], '.')
        ylim(1,-1)
        grid()


        ii = argsort(pr[:,6])
        a,b=polyfit(pr[:,6],pr[:,1]-pr[:,2],1)
        plot(pr[ii,6], a*pr[ii,6]+b,'-g')
        a,b,c,d=polyfit(pr[:,6],pr[:,1]-pr[:,2],3)
        plot(pr[ii,6], a*pr[ii,6]**3+b*pr[ii,6]**2+c*pr[ii,6]+d,'-r')


        # #xx = mgrid[60:160:.1] # img 9 trecho
        # xx = mgrid[100:260:.1] # img 9 trecho
        # #plot(xx, a*xx+b)
        # plot(xx, a*xx**3+b*xx**2+c*xx+d)

